#!/usr/bin/env python
# -*- coding: UTF-8 -*-

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LinearRegression
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 读取数据
data = pd.read_excel('../data/恒转速数据.xlsx')

# 提取输入特征和目标变量
X = data[['Speed', 'Power']]
Y = data[['Flow', 'Head']]

# 数据划分
X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size=0.2, random_state=42)

# 多项式特征转换
poly = PolynomialFeatures(degree=2)  # 可以调整degree值来获得不同的多项式阶数
X_train_poly = poly.fit_transform(X_train)
X_test_poly = poly.transform(X_test)

# 训练线性回归模型
model = LinearRegression()
model.fit(X_train_poly, y_train)

# 生成网格数据以适应三维曲面
speed_range = np.linspace(X['Speed'].min(), X['Speed'].max(), 50)
power_range = np.linspace(X['Power'].min(), X['Power'].max(), 50)
speed_mesh, power_mesh = np.meshgrid(speed_range, power_range)

# 转换为多项式特征
X_mesh = poly.transform(np.column_stack((speed_mesh.ravel(), power_mesh.ravel())))

# 预测流量和扬程
y_mesh_pred = model.predict(X_mesh)

# 绘制三维曲面图（假设绘制流量）
fig = plt.figure(figsize=(14, 7))

# 三维曲面图
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(speed_mesh, power_mesh, y_mesh_pred[:, 0].reshape(speed_mesh.shape),
                cmap='viridis', alpha=0.8)

# 设置标签
ax.set_xlabel('转速 (Speed)')
ax.set_ylabel('功率 (Power)')
ax.set_zlabel('流量 (Flow)')
ax.set_title('转速与功率对流量的影响')

plt.show()

